#########################################################################
# Time-stamp: <05/07/18 15:15:09 ostolop>
#########################################################################
package EP::EPC::Core::Advanced_Clustering_Comparison;

use strict;

use EP::Config;
use EP::Common;
use EP::Common::General;

use EP::EPC::Core;
use EP::EPC::Core::Component;
use EP::EPC::Core::Clustering;

use Text::Template;

use EP::Visualization::Grid;
use EP::Visualization::Hooks;
use EP::Visualization::HeatmapGraphic;
use EP::Visualization::LineGraphic;
use EP::Visualization::TreeGraphic;
use EP::Visualization::AnnotationGraphic;
use EP::Visualization::ClusteringComparisonGraphic;

use PDL;
use PDL::NiceSlice;
use PDL::NetCDF;

our @ISA = qw ( EP::EPC::Core::Component );

sub process_component {
  my $self = shift;

  my ($result, $svg_file, $width, $height) = $self->run_clustering_comparison ();
  my $process_result = {
			acc_svg  => { type => "image/svg+xml",
				      label => "Advanced Clustering Comparison Visualization",
				      value => "$EP::Config::h->{HTML_HREF}/$EP::Config::h->{WWW_TMP}/$svg_file" },
			acc_svg_width => $width,
			acc_svg_height => $height,
			src_dataset_id => $self->{session}->current_root_dataset->id,
			dst_dataset_id => $self->{session}->current_root_dataset->id
		       };

  $self->{CD}->{process_result} = $process_result;
  return $self->record_analysis_history ();
}

sub run_clustering_comparison {
  my $self = shift;

  my ( $clust_action_id1, $clust_action_id2 ) = split /,/, $self->{query}->param ( "selected_actions" );

  my $clust_action1 = retrieve EP::Common::DBI::Action ( $clust_action_id1 ) if $clust_action_id1;
  my $clust_action2 = retrieve EP::Common::DBI::Action ( $clust_action_id2 ) if $clust_action_id2;

  # test whether (if both were found) they were from the same dataset
  if ( defined $clust_action1 and defined $clust_action2 ) {
    die "The two clusterings should originate from the same dataset!" if $clust_action1->parent_action->dst_dataset != $clust_action2->parent_action->dst_dataset;
  }

  # if only one clustering action was selected, run a clustering action of diff't type on the same dataset as the clustered one
  if ( $clust_action1 and not $clust_action2 ) {
    $clust_action2 = $self->run_tree_clustering ( $clust_action1->parent_action->dst_dataset ) if $clust_action1->component->name eq 'K-means/K-medoids Clustering';
    $clust_action2 = $self->run_flat_clustering ( $clust_action1->parent_action->dst_dataset ) if $clust_action1->component->name eq 'Hierarchical Clustering';
  }

  if ( not $clust_action1 and not $clust_action2 ) {
    $clust_action1 = $self->run_tree_clustering ( $self->{src_dataset} );
    $clust_action2 = $self->run_flat_clustering ( $self->{src_dataset} );
  }

  my ($tree, $tree_clustering_cdf)
                          = $self->format_tree_clustering ( $clust_action1->component->name eq 'Hierarchical Clustering' ? $clust_action1 : $clust_action2 );
  my $flat_clustering_cdf = $self->format_flat_clustering ( $clust_action1->component->name eq 'Hierarchical Clustering' ? $clust_action2 : $clust_action1 );

  my $comparison_cdf = $self->compare_clusterings ( $tree_clustering_cdf, $flat_clustering_cdf,
						    $self->{query}->param ( "look_ahead" ),
						    $self->{query}->param ( "scoring_fn" ),
						    $self->{query}->param ( "cte" ) );

  return $self->create_visualization ( $clust_action1, $clust_action2, $tree, $tree_clustering_cdf, $flat_clustering_cdf, $comparison_cdf );
}

# actually execute hierarchical clustering
sub run_tree_clustering {
  my $self = shift;
  my ( $dataset ) = @_;

  throw EP::Common::Error ( -value => AFFY_TRANSFORM_RAW_DATA ) if $dataset->type eq "affy";

  my $clustering_epc = new EP::EPC::Core::Clustering ( $self->{query}, $self->{CD}, 3 );
  $clustering_epc->{src_dataset} = $dataset;
  my $clustering_params = new CGI ( { clustering_action => 'hierarchical',
				      distance          => $self->{query}->param("hier_distance"),
				      algorithm         => $self->{query}->param("hier_algorithm"),
				      impute_row_avg    => $self->{query}->param("impute_row_avg")
				    } );

  my $current_action = $self->{session}->current_action;
  my $current_root_dataset = $self->{session}->current_root_dataset;

  my $new_action = $clustering_epc->process_component ( $clustering_params );

  if ( $new_action ) {
    # reset current action to source dataset
    $self->{session}->current_action ( $current_action );
    $self->{session}->current_root_dataset ( $current_root_dataset );
    $self->{CD}->{process_result}->{current_action_id} = $current_action;
    return $new_action;
  }

  die "Failed to run hierarchical clustering!";
}

# actually execute flat clustering
sub run_flat_clustering {
  my $self = shift;
  my ( $dataset ) = @_;

  throw EP::Common::Error ( -value => AFFY_TRANSFORM_RAW_DATA ) if $dataset->type eq "affy";

  my $clustering_epc = new EP::EPC::Core::Clustering ( $self->{query}, $self->{CD}, 4 );
  $clustering_epc->{src_dataset} = $dataset;
  my $clustering_params = new CGI ( { clustering_action => $self->{query}->param("flat_clustering_action"),
				      K                 => $self->{query}->param("K"),
				      distance          => $self->{query}->param("flat_distance"),
				      start             => $self->{query}->param("start"),
				      impute_row_avg    => $self->{query}->param("impute_row_avg")
				    } );

  my $current_action = $self->{session}->current_action;
  my $current_root_dataset = $self->{session}->current_root_dataset;

  my $new_action = $clustering_epc->process_component ( $clustering_params );

  if ( $new_action ) {
    # reset current action to source dataset
    $self->{session}->current_action ( $current_action );
    $self->{session}->current_root_dataset ( $current_root_dataset );
    $self->{CD}->{process_result}->{current_action_id} = $current_action;
    return $new_action;
  }

  die "Failed to run flat clustering!";
}

sub create_visualization {
  my $self = shift;
  my ( $clust_action1, $clust_action2, $tree, $tree_clustering_cdf, $flat_clustering_cdf, $comparison_cdf ) = @_;

  my $tree_clustering_action = $clust_action1->component->name eq 'Hierarchical Clustering' ? $clust_action1 : $clust_action2;
  my $flat_clustering_action = $clust_action1->component->name eq 'Hierarchical Clustering' ? $clust_action2 : $clust_action1;

  my $tree_clustering_vis;
  my $cell_width  = 13;

  my $raw_data = $tree_clustering_action->dst_dataset->data->raw_data;
  my $num_clusters = $comparison_cdf->get('cluster_order')->dim(0);
  my $num_objects = $raw_data->dim(1);

  my $plot_height = 100;
  my $cell_height = $plot_height * $num_clusters / $num_objects;

  $tree_clustering_vis = new EP::Visualization::TreeGraphic ( data => $tree, leaf_height => $cell_height, reorder => $comparison_cdf->get('leaves_order'), cutpoints => [list $comparison_cdf->get('tree_cutpoints')], interactive => 0 );
  $tree_clustering_vis->setOrientation ( -90 );

  my $heatmap = new EP::Visualization::HeatmapGraphic ( data => $raw_data,
							cell_height => $cell_height, cell_width => $cell_width
						      );
  $heatmap->orderVerticalBy ( $tree_clustering_vis );

  my $kmeans_clustering_grid = new EP::Visualization::Grid ( 2, 2 * $comparison_cdf->get('cluster_order')->dim(0) );
  my $kmeans_cell_array;

  my $flat_clustering = new PDL::NetCDF ( $flat_clustering_cdf );
  my $corr = $flat_clustering->get ('correspondence');

  my @v_hooks = sort { $a <=> $b} keys %{$tree_clustering_action->dst_dataset->data->{row}};
  my @h_hooks = sort { $a <=> $b} keys %{$tree_clustering_action->dst_dataset->data->{column}};
  
  my $V_hook = new EP::Visualization::Hooks ( \@v_hooks );
  my $H_hook = new EP::Visualization::Hooks ( \@h_hooks );

  foreach ( list $comparison_cdf->get('cluster_order') ) {
    my $cluster_members = $corr ( which ( $corr ( , 1 ) == $_ ) , 0 )->clump(-1);
    my $d = $raw_data ( :, $cluster_members - 1 );
    my $cluster_lineplot = new EP::Visualization::LineGraphic ( data => $d, plot_height => $plot_height - 30, cell_width => $cell_width, interactive => 0 );
    $cluster_lineplot->setVerticalHookSet ( $V_hook );

    push @$kmeans_cell_array, [ undef, $cluster_lineplot ];
    push @$kmeans_cell_array, [ undef, undef ];
 }
  $kmeans_clustering_grid->fillCellArray ( $kmeans_cell_array );

  my $clustering_comparison_vis = new EP::Visualization::ClusteringComparisonGraphic ( tree => $tree_clustering_vis,
										       heatmap => $heatmap,
										       cell_height => $cell_height,
										       comparison => $comparison_cdf );

  $heatmap->setVerticalHookSet   ( $V_hook );
  $heatmap->setHorizontalHookSet ( $H_hook );

  #create annotation
  my @row_annots;
  push @row_annots, $tree_clustering_action->dst_dataset->data->{row}->{$_} foreach @v_hooks;

  my $row_annotation = new EP::Visualization::AnnotationGraphic ( data => \@row_annots, 'font-size' => $cell_height, 'no-color-efs' => 1 );
  $row_annotation->setHook ( $V_hook );
  $row_annotation->orderBy    ( $tree_clustering_vis );

  my @col_annots;
  push @col_annots, $tree_clustering_action->dst_dataset->data->{column}->{$_} foreach @h_hooks;

  my $col_annotation = new EP::Visualization::AnnotationGraphic ( data => \@col_annots, 'font-size' => $cell_width );
  $col_annotation->setHook ( $H_hook );
  $col_annotation->setOrientation(-90);

  my $top_grid = new EP::Visualization::Grid ( 6, 2 );
  $top_grid->fillCellArray ( [ [ undef, $col_annotation, undef, undef, undef, undef ],
			       [ $tree_clustering_vis, $heatmap, $row_annotation, $clustering_comparison_vis, $kmeans_clustering_grid, undef ] ] );

  $top_grid->to_svg ( toFile => "$EP::Config::h->{TMP}/cc.$$.svgz", width => $top_grid->getGridWidth(1), height => $top_grid->getGridHeight(1) );
  return ( "", "cc.$$.svgz",
	   $top_grid->getGridWidth(1),
	   $top_grid->getGridHeight(1) );
}

sub compare_clusterings {
  my $self = shift;
  my ( $tree_clustering_cdf, $flat_clustering_cdf, $look_ahead, $scoring_fn, $cte ) = @_;

  my $output_cdf = new File::Temp ( UNLINK => 0, SUFFIX => '.nc' );

  my $tmpl_values = { html_root           => $EP::Config::h->{HTML_ROOT},
		      tree_clustering_cdf => $tree_clustering_cdf,
                      flat_clustering_cdf => $flat_clustering_cdf,
		      look_ahead          => $look_ahead,
		      scoring_fn          => $scoring_fn,
		      cte                 => $cte,
		      output_cdf          => "$output_cdf"
		    };

  $self->{session}->status ( "Running clustering comparison ..." );
  my $Rout = launch_R ( "$EP::Config::h->{TMPL}/R/ep.run_clustering_comparison.tmpl.R", $tmpl_values );

  if ( defined EP::Common::Error->prior() and EP::Common::Error->prior()->value() == SYS_CALL_WITH_NONZERO_STATUS ) {
    die $Rout;# TODO: Error!
  }

  return new PDL::NetCDF ( "$output_cdf" );
}

sub format_flat_clustering {
  my $self = shift;

  my ($clustering_action) = @_;

  my ($clustering_filename) = EP::Common::DBI::Action_Metadata->search ( { action => $clustering_action->id, name => "clustering" } );
  my $clustering_file = $EP::Config::h->{EPNG_USRDATA_PATH} . "/" . $clustering_action->dst_dataset->folder->location . "/" . $clustering_filename->value;

  open ( CLUSTERING,
	 $clustering_file )
    or die "Couldn't open clustering file $clustering_file: $!";

  return "$clustering_file.nc" if -e "$clustering_file.nc";

  my @CC;
  while ( <CLUSTERING> ) {
    chomp();
    if(/^O:/) {
      my @object = split(/\t/, $');
      push @CC, [$object[0], $object[1]];
    }
  }
  close ( CLUSTERING );

  my $cc_cdf = new PDL::NetCDF ( ">$clustering_file.nc" );
  $cc_cdf->put ( 'correspondence', ['r','c'], transpose pdl \@CC );
  $cc_cdf->close();

  return "$clustering_file.nc";
}

# Sub: format_tree_clustering
#   Reads in Jaak's clustering format, outputs a file close to R's native format (hclust)
sub format_tree_clustering {
  my $self = shift;

  my ($clustering_action) = @_;

  my ($clustering_filename) = EP::Common::DBI::Action_Metadata->search ( { action => $clustering_action->id, name => "rows_clust" } );
  my $clustering_file = $EP::Config::h->{EPNG_USRDATA_PATH} . "/" . $clustering_action->dst_dataset->folder->location . "/" . $clustering_filename->value;

  #return ($clustering_file, "$clustering_file.nc") if -e "$clustering_file.nc";

  my $tree = {};
  read_in_tree ( $tree, $clustering_file );
  my $hclust = to_hclust ( $tree );

  my $hclust_cdf = new PDL::NetCDF ( ">$clustering_file.nc" );

  $hclust_cdf->put ( 'merge',  ['r','n'], transpose pdl $hclust->{merge} );
  $hclust_cdf->put ( 'height', ['h'],     pdl $hclust->{height} );
  $hclust_cdf->put ( 'order',  ['o'],     pdl $hclust->{order} );

  $hclust_cdf->close();

  return ($tree, "$clustering_file.nc");
}

sub run_clusterings {
  my $self = shift;

  if (  $self->{query}->param ( "selected_actions" ) ) {
    my $clust;
  }

  EP::EPC::Core::Clustering::exp::process_clustering ( $self, $self->{src_dataset}, $self->{query} );
}

sub create_clustering_correspondence {
  my ( $clustering_file ) = @_;
  my $clustering_correspondence = $clustering_file . ".correspondence.$$.tab";

  open ( CLUSTERING, $clustering_file ) or
    die "Couldn't open clustering file $clustering_file: $!";
  open ( CLUSTERING_CORRESPONDENCE,">$clustering_correspondence") or
    die "Couldn't create output correspondence file $clustering_correspondence: $!";

  my @object;
  while ( <CLUSTERING> ) {
    chomp();
    if(/^O:/) {
      @object = split(/\t/, $');
      print CLUSTERING_CORRESPONDENCE "$object[0]\t$object[1]\n";
    }
  }

  close(CLUSTERING_CORRESPONDENCE);
  close(CLUSTERING);

  return $clustering_correspondence;
}

sub to_hclust {
  my $self = shift;

  my @RTREE;
  foreach ( sort { $self->{DIST}{$a} <=> $self->{DIST}{$b} } keys %{$self->{DIST}} ) {
    my ($l, $r) = split /:/, $_;
    my $m = $self->{SUBTREE}{$_};

    my ($nl, $nr) = ( -$l, -$r );

    if ( $self->{TREE_L}->[$self->{TREE_L}->[$m]] != 0 ) {
      foreach ( my $i = $#RTREE; $i >= 0; $i-- ) {
	$nl = $i + 1 and last if $l == $RTREE[$i]->[0] or $l == $RTREE[$i]->[1];
      }
    }

    if ( $self->{TREE_L}->[$self->{TREE_R}->[$m]] != 0 ) {
      foreach ( my $i = $#RTREE; $i >= 0; $i-- ) {
	$nr = $i + 1 and last if $r == $RTREE[$i]->[1] or $r == $RTREE[$i]->[0];
      }
    }

    push @RTREE, [$l, $r, $nl, $nr, $self->{DIST}{$_}];
  }

  my @merge;
  push @merge, [ @$_[2,3] ] foreach @RTREE;

  return { merge  => [ @merge ],
	   height => [ sort { $a <=> $b } values %{$self->{DIST}} ],
	   order  => [ map { $_ + 1 } treewalk ( $self ) ] };
}

sub to_phylip {
  my $self = shift;
  my ($N) = @_;

  $N = 0 if not $N;

  return "$self->{TREE_R}->[$N]:1" if $self->{TREE_L}->[$N] == 0;

  my $l = $self->{ORIG_TREE}->[$self->{TREE_L}->[$N]];
  my $r = $self->{ORIG_TREE}->[$self->{TREE_R}->[$N]];
  my $DIST = $self->{DIST}->{"$l:$r"}; # Distance between left and right subtrees

  return "(" .
    to_phylip ( $self, $self->{TREE_L}->[$N] )
      . "," .
    to_phylip ( $self, $self->{TREE_R}->[$N] ) .
      "):$DIST";
}

sub read_tree_data_from_file {
  my ($file) = shift;

  my @DATA;

  open ( CLUSTER, $file ) or die "Couldn't open the hierarchical clustering file $file: $!\n";
  while (<CLUSTER>) {
    next if /^\s*#/;

    # From the original clustering output
    my @data_row = split /\s+/;
    push @DATA, \@data_row;
  }
  close CLUSTER;

  return \@DATA;
}

sub read_in_tree {
  my $self = shift;
  my ( $tree_data ) = @_;

  my %LL;
  my @TREEDOWN;
  my $TREENODE = 0;

  my (@TREE_L, @TREE_R, %DIST, %SUBTREE, @ORIG_TREE);

  my $DATA;
  $DATA = read_tree_data_from_file ( $tree_data ) if -f $tree_data;
  $DATA = $tree_data if ref ( $tree_data ) eq "ARRAY";

  die "Couldn't read in tree data!" if not $DATA;

  foreach my $data_row ( @{$DATA} ) {
    # From the original clustering output
    my ( $LH, $RH, $LHS, $RHS, $D ) = @$data_row;

    $TREEDOWN[1] = [ 0, $LHS + $RHS ] if not exists $TREEDOWN[1];

    # @TREEDOWN  - how much space was reserved for subtree starting with that node
    my ( $TREENODE, $count ) = @{$TREEDOWN[$LH]};

    # $LHTREE, $RHTREE, $TREENODE -  tree-node id's on tree structure
    my $LHTREE = $TREENODE + 1;
    my $RHTREE = $TREENODE + nodes_in_subtree($LHS) + 1;

    ( $TREE_L   [$TREENODE], $TREE_R   [$TREENODE] ) = ( $LHTREE,           $RHTREE);
    ( $ORIG_TREE[$LHTREE  ], $ORIG_TREE[$RHTREE  ] ) = ( $LH,               $RH );
    ( $TREEDOWN [$LH      ], $TREEDOWN [$RH      ] ) = ( [ $LHTREE, $LHS ], [ $RHTREE, $RHS ] );

    $DIST{"$LH:$RH"}    = $D;
    $SUBTREE{"$LH:$RH"} = $TREENODE;

    # left leaf
    ( $TREE_L[$LHTREE], $TREE_R[$LHTREE] ) = ( 0, $LH ) if $LHS == 1;

    # right leaf
    ( $TREE_L[$RHTREE], $TREE_R[$RHTREE] ) = ( 0, $RH ) if $RHS == 1;
  }

  $self->{nr_leaves} = $$DATA[0][2] + $$DATA[0][3];
  $self->{TREE_L} = \@TREE_L;
  $self->{TREE_R} = \@TREE_R;
  $self->{DIST}   = \%DIST;
  $self->{SUBTREE} = \%SUBTREE;
  $self->{ORIG_TREE} = \@ORIG_TREE;

  $self->{data} = $DATA;
}

sub nodes_in_subtree {
  my ( $leaves ) = $_[0];
  return $leaves * 2 - 1;
}

sub treewalk {
  my $self = shift;
  my ($N) = @_;

  $N = 0 if not $N;

  if ( $self->{TREE_L}->[$N] == 0 ) {
    my $leaf = $self->{TREE_R}->[$N];
    return $leaf - 1;		# subtract 1 because indices should be zero-based
  }

  return ( treewalk ( $self, $self->{TREE_L}->[$N] ),
	   treewalk ( $self, $self->{TREE_R}->[$N] ) );
}

sub left_treewalk {
  my $self = shift;
  my ($N) = @_;

  $N = 0 if not $N;

  if ( $self->{TREE_L}->[$N] == 0 ) {
    my $leaf = $self->{TREE_R}->[$N];
    return $leaf - 1;		# subtract 1 because indices should be zero-based
  }

  return ( $self->left_treewalk ( $self->{TREE_L}->[$N] ) );
}

sub right_treewalk {
  my $self = shift;
  my ($N) = @_;

  $N = 0 if not $N;

  if ( $self->{TREE_L}->[$N] == 0 ) {
    my $leaf = $self->{TREE_R}->[$N];
    return $leaf - 1;		# subtract 1 because indices should be zero-based
  }

  return ( $self->right_treewalk ( $self->{TREE_R}->[$N] ) );
}

sub svg_html {
  my ($svg_file, $width, $height) = @_;
  my $svg_html = <<__SVG_APPLET__;
    <EMBED ALIGN="center" SRC="$EP::Config::h->{HTML_HREF}/$EP::Config::h->{WWW_TMP}/$svg_file" NAME="SVGEmbed" HEIGHT="$height" WIDTH="$width" TYPE="image/svg-xml" PLUGINSPAGE="http://www.adobe.com/svg/viewer/install/">
__SVG_APPLET__

  return $svg_html;
}

1;
